package com.zxy.leetcode._00200_00299._00210_00219;

/**
 * https://leetcode-cn.com/problems/kth-largest-element-in-an-array/
 *
 * 数组中的第K个最大元素
 *
 * 在未排序的数组中找到第 k 个最大的元素。
 * 请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。
 *
 * 标签：分治，快排
 */
public class Test00215 {

    public static void main(String[] args) {
        Test00215 test = new Test00215();
        int[] nums1 = {3,2,1,5,6,4};
        System.out.println(test.findKthLargest(nums1, 2) == 5);

        int[] nums2 = {3,2,3,1,2,4,5,5,6};
        System.out.println(test.findKthLargest(nums2, 4) == 4);

        int[] nums3 = {1, 2, 3, 4, 5, 6};
        System.out.println(test.findKthLargest(nums3, 1) == 6);
        System.out.println(test.findKthLargest(nums3, 2) == 5);

        int[] nums4 = {2, 1};
        System.out.println(test.findKthLargest(nums4, 2) == 1);
    }


    public int findKthLargest(int[] nums, int k) {
        return partition(nums, k, 0, nums.length - 1);
    }

    public int partition(int[] nums, int k, int left, int right) {
        // 极端情况下，nums本身可能基本有序的，导致下面选择的中心轴无效
        // 标准的做法是用random
        swap(nums, left, (left + right)/2);

        int pivot = nums[left]; // 中心轴
        int i = left;
        int j = right;
        while (i < j) {
            while (i < j && nums[j] >= pivot) j --;

            swap(nums, i, j);

            while (i < j && nums[i] <= pivot) i ++;

            swap(nums, i, j);
        }

        // 此时 i==j
        int pivotIndex = i;
        if (right - pivotIndex == k - 1) {
            return nums[pivotIndex];
        }

        if (right - pivotIndex < k - 1) {
            return partition(nums, k - (right - pivotIndex + 1), left, pivotIndex - 1);
        }

        return partition(nums, k, pivotIndex + 1, right);

    }

    private void swap(int[] nums, int a, int b) {
        if (a == b) {
            return;
        }
        int tmp = nums[a];
        nums[a] = nums[b];
        nums[b] = tmp;
    }

}
